import os
import json
import numpy as np
from PIL import Image
from cyw_devkit.core.io import files_in_fir, read_points
from cyw_devkit.core.transform import SuperTransform
from cyw_devkit.core.visualizer import draw_pts2img, just_look_array
from cyw_devkit.core.zmath import to_matrix4x4
from cyw_devkit.core.dataset import Data, BBoxes


class BaseFiles:
    def __init__(self, father_path, dataset_name, data_dir):
        self.globel_path = os.path.join(father_path, dataset_name, data_dir)
        self.local_path = os.path.join(dataset_name, data_dir)
        self.files = files_in_fir(self.globel_path)

    def get_iterm(self, index):
        if index < len(self.files):
            return os.path.join(self.local_path, self.files[index])
        else:
            return None



def get_stamp(path):
    return int(os.path.basename(path).split('.')[0])


class CywDataset():
    '''
    datasets:
        dataset0:
            calib:
                calib.json
            lidar_bin:
            radar_bin:
            camera:
                camera_75
    '''

    def __init__(self, dataset_path, lidar_dim=4, radar_dim=5, normal_img=True, motion=False):
        '''
        cyw radar_dim=5 carla radar_dim=7
        '''
        assert os.path.exists(dataset_path), f"{dataset_path} does not exist!"
        self.father_path = os.path.dirname(dataset_path)
        self.dataset_name = os.path.basename(dataset_path)
        self.lidar_dim = lidar_dim
        self.radar_dim = radar_dim
        self.normal_img = normal_img
        self.motion = motion
        self.parse_path()

    @property
    def dataset_path(self):
        return os.path.join(self.father_path, self.dataset_name)

    def parse_path(self):
        super_transform = SuperTransform(
            local_csv=os.path.join(self.dataset_path, 'localization', 'localization.csv'),
            calib_json=os.path.join(self.dataset_path, 'calib', 'calib.json'))
        with open(os.path.join(self.dataset_path, 'calib', 'calib.json'), 'r') as f:
            calib_info = json.load(f)

        lidar_files = BaseFiles(father_path=self.father_path, dataset_name=self.dataset_name, data_dir='lidar_bin')
        radar_files = BaseFiles(father_path=self.father_path, dataset_name=self.dataset_name, data_dir='radar_bin')
        if not self.motion:
            label_files = BaseFiles(father_path=self.father_path, dataset_name=self.dataset_name, data_dir='samples/label')
        else:
            label_files = BaseFiles(father_path=self.father_path, dataset_name=self.dataset_name, data_dir='samples/label_e2e')

        camera_files = dict()
        if os.path.exists(os.path.join(self.dataset_path, 'camera')):
            for dir in os.listdir(os.path.join(self.dataset_path, 'camera')):
                if (self.normal_img and dir[-6:] == 'normal') or (not self.normal_img and dir[-6:] != 'normal'):
                    camera_files[dir[:8]] = BaseFiles(father_path=self.father_path, dataset_name=self.dataset_name,
                                                      data_dir=os.path.join('camera', dir))
        frame_size = len(os.listdir(os.path.join(self.father_path, self.dataset_name, 'lidar_bin')))
        samples_dir = os.path.join(self.dataset_path, 'samples', 'samples.json')
        with open(samples_dir, 'r') as f:
            json_samples = json.load(f)

        self.data_dict = {
            'super_transform': super_transform,
            'calib_info': calib_info,
            'frames': [],
            'key_frames': json_samples['key_index'],
        }
        for idx in range(frame_size):
            frame = {'idx': idx,
                     'lidar': lidar_files.get_iterm(index=idx),
                     'radar': radar_files.get_iterm(index=idx),
                     'camera': dict(),
                     'label': label_files.get_iterm(index=self.data_dict['key_frames'].index(idx)) if idx in self.data_dict['key_frames'] else None,
                     }
            for camera_id, camera_file in camera_files.items():
                frame['camera'][camera_id] = camera_file.get_iterm(index=idx)
            self.data_dict['frames'].append(frame)
        pass

    def get_label(self, index):
        path = self.data_dict['frames'][index]['label']
        if path is None: # 非关键帧返回None
            return None
        with open(os.path.join(self.father_path, path), 'r') as f:
            label = json.load(f)
        boxes = BBoxes(label)
        data = Data(stamp=get_stamp(path), frame_id='base_link', data=boxes)
        return data

    def get_lidar(self, index):
        path = self.data_dict['frames'][index]['lidar']
        pcs = read_points(os.path.join(self.father_path, path), dim=self.lidar_dim)
        data = Data(stamp=get_stamp(path), frame_id='base_link', data=pcs)
        return data

    def get_radar(self, index):
        path = self.data_dict['frames'][index]['radar']
        pcs = read_points(os.path.join(self.father_path, path), dim=self.radar_dim)
        data = Data(stamp=get_stamp(path), frame_id='base_link', data=pcs)
        return data

    def get_camera(self, index, camera_id):
        path = self.data_dict['frames'][index]['camera'][camera_id]
        img = Image.open(os.path.join(self.father_path, path))
        data = Data(stamp=get_stamp(path), frame_id=camera_id, data=img)
        return data

    def get_all_camera(self, index):
        camera_dict = dict()
        for camera_id, camera_file in self.data_dict['frames'][index]['camera'].items():
            camera_dict[camera_id] = self.get_camera(index, camera_id)
        return camera_dict

    def get_camera_mat(self, camera_id):
        assert camera_id in self.data_dict['calib_info']
        if self.normal_img:
            mat33 = self.data_dict['calib_info'][camera_id]['normal']['K']
        else:
            mat33 = np.array(self.data_dict['calib_info'][camera_id]['ori']['K'])
        return to_matrix4x4(mat33)

    @property
    def keyframe_idxs(self):
        return self.data_dict['key_frames']

    def get_frame_path(self, index):
        return self.data_dict['frames'][index]

    def get_frame_stamp(self, index):
        frame_path = self.get_frame_path(index)
        return get_stamp(frame_path['lidar'])

    def __getitem__(self, index):
        lidar = self.get_lidar(index)
        radar = self.get_radar(index)
        camera_dict = self.get_all_camera(index)
        label = self.get_label(index)
        frame_data = {
            'lidar': lidar,
            'radar': radar,
            'camera_dict': camera_dict,
            'label': label
        }
        return frame_data

    def __len__(self):
        return len(self.data_dict['frames'])


if __name__ == '__main__':
    cyw = CywDataset(dataset_path='/media/adt/ZWH4T/ZWH/bags/dataset/sll/test_out/__roate')
    frame_data = cyw.__getitem__(0)

    for key_idx in cyw.keyframe_idxs:
        key_frame = cyw.__getitem__(key_idx)
        pass

    camera2base = cyw.data_dict['super_transform'].get_tf_from_space(source_point='camera75', target_point='base_link')
    camera2img = cyw.get_camera_mat('camera75')
    lidar = frame_data['lidar']
    image = frame_data['camera_dict']['camera75']
    image_new = draw_pts2img(points=lidar.data, image=np.array(image.data),
                             tfs=[np.linalg.inv(camera2base), camera2img])
    just_look_array(image_new)
    pass
